import logging
from typing import List, Dict, Any, Optional
from langchain.schema import Document
import numpy as np
from config.thresholds import RAGThresholds
import os
import requests
import json
from config import Config

logger = logging.getLogger(__name__)

class Reranker:
    """重排序器 - 使用方舟API进行重排序"""
    
    def __init__(self, model_name: str = None):
        """
        初始化重排序器
        
        Args:
            model_name: 重排序模型名称，默认使用方舟的重排序模型
        """
        self.model_name = model_name or "doubao-seed-1-6-250615"
        self.api_key = Config.OPENAI_API_KEY
        self.api_base = Config.OPENAI_BASE_URL
        
        logger.info(f"初始化重排序器，使用模型: {self.model_name}")
    
    def _get_rerank_scores(self, query: str, documents: List[str]) -> List[float]:
        """使用方舟模型获取重排序分数"""
        try:
            headers = {
                "Authorization": f"Bearer {self.api_key}",
                "Content-Type": "application/json"
            }
            
            # 构建重排序请求
            data = {
                "model": self.model_name,
                "messages": [
                    {
                        "role": "system",
                        "content": "你是一个专业的文档相关性评估助手。请根据用户查询，对提供的文档进行相关性评分。评分范围0-1，1表示最相关，0表示不相关。请只返回数字分数，不要其他内容。"
                    },
                    {
                        "role": "user",
                        "content": f"查询：{query}\n\n请对以下文档进行相关性评分：\n" + "\n".join([f"{i+1}. {doc[:200]}..." for i, doc in enumerate(documents)])
                    }
                ],
                "temperature": 0.1,
                "max_tokens": 100
            }
            
            response = requests.post(
                f"{self.api_base}/chat/completions",
                headers=headers,
                json=data,
                timeout=30
            )
            
            if response.status_code == 200:
                result = response.json()
                content = result["choices"][0]["message"]["content"]
                
                # 解析分数
                try:
                    # 尝试提取数字
                    import re
                    scores = re.findall(r'\d+\.?\d*', content)
                    if len(scores) >= len(documents):
                        return [float(score) for score in scores[:len(documents)]]
                    else:
                        # 如果解析失败，返回默认分数
                        return [0.5] * len(documents)
                except:
                    return [0.5] * len(documents)
            else:
                logger.error(f"重排序API调用失败: {response.status_code} - {response.text}")
                return [0.5] * len(documents)
                
        except Exception as e:
            logger.error(f"重排序API调用异常: {str(e)}")
            return [0.5] * len(documents)
    
    def rerank_documents(
        self,
        query: str,
        documents: List[Document],
        top_k: int = RAGThresholds.DEFAULT_TOP_K,
        threshold: float = RAGThresholds.DEFAULT_RERANK_THRESHOLD
    ) -> List[Document]:
        """
        对文档进行重排序
        
        Args:
            query: 查询文本
            documents: 待重排序的文档列表
            top_k: 返回前k个文档
            threshold: 相关性阈值，低于此值的文档将被过滤
        
        Returns:
            重排序后的文档列表
        """
        try:
            if not documents:
                return []
            
            logger.info(f"开始重排序，查询: {query[:50]}...，文档数: {len(documents)}")
            
            # 获取文档内容
            doc_contents = [doc.page_content for doc in documents]
            
            # 使用方舟模型获取重排序分数
            scores = self._get_rerank_scores(query, doc_contents)
            
            # 将分数与文档配对
            doc_scores = list(zip(documents, scores))
            
            # 按分数降序排序
            doc_scores.sort(key=lambda x: x[1], reverse=True)
            
            # 过滤低相关性文档
            filtered_docs = [doc for doc, score in doc_scores if score >= threshold]
            
            # 返回前top_k个文档
            reranked_docs = filtered_docs[:top_k]
            
            # 添加重排序信息到元数据
            for i, (doc, score) in enumerate(doc_scores[:top_k]):
                if doc in reranked_docs:
                    doc.metadata["rerank_score"] = float(score)
                    doc.metadata["rerank_rank"] = i + 1
            
            logger.info(f"重排序完成，返回 {len(reranked_docs)} 个文档")
            return reranked_docs
            
        except Exception as e:
            logger.error(f"重排序失败: {str(e)}")
            # 出错时返回原始排序
            return documents[:top_k]
    
    def rerank_with_metadata(
        self,
        query: str,
        documents: List[Document],
        top_k: int = RAGThresholds.DEFAULT_TOP_K,
        threshold: float = RAGThresholds.DEFAULT_RERANK_THRESHOLD,
        use_metadata: bool = True
    ) -> Dict[str, Any]:
        """
        重排序并返回详细信息
        
        Args:
            query: 查询文本
            documents: 待重排序的文档列表
            top_k: 返回前k个文档
            threshold: 相关性阈值
            use_metadata: 是否使用元数据增强排序
        
        Returns:
            包含重排序结果的字典
        """
        try:
            if not documents:
                return {
                    "reranked_documents": [],
                    "scores": [],
                    "filtered_count": 0,
                    "total_count": 0
                }
            
            # 基础重排序
            reranked_docs = self.rerank_documents(query, documents, top_k, threshold)
            
            # 收集分数信息
            scores = []
            for doc in reranked_docs:
                score = doc.metadata.get("rerank_score", 0.0)
                scores.append(score)
            
            # 统计信息
            filtered_count = len(reranked_docs)
            total_count = len(documents)
            
            result = {
                "reranked_documents": reranked_docs,
                "scores": scores,
                "filtered_count": filtered_count,
                "total_count": total_count,
                "average_score": np.mean(scores) if scores else 0.0,
                "max_score": max(scores) if scores else 0.0,
                "min_score": min(scores) if scores else 0.0
            }
            
            logger.info(f"重排序统计: 过滤前 {total_count} 个，过滤后 {filtered_count} 个，平均分数 {result['average_score']:.3f}")
            
            return result
            
        except Exception as e:
            logger.error(f"重排序失败: {str(e)}")
            return {
                "reranked_documents": documents[:top_k],
                "scores": [0.0] * min(len(documents), top_k),
                "filtered_count": min(len(documents), top_k),
                "total_count": len(documents),
                "error": str(e)
            }
    
    def batch_rerank(
        self,
        queries: List[str],
        documents_list: List[List[Document]],
        top_k: int = RAGThresholds.DEFAULT_TOP_K,
        threshold: float = RAGThresholds.DEFAULT_RERANK_THRESHOLD
    ) -> List[List[Document]]:
        """
        批量重排序
        
        Args:
            queries: 查询列表
            documents_list: 每个查询对应的文档列表
            top_k: 返回前k个文档
            threshold: 相关性阈值
        
        Returns:
            重排序后的文档列表
        """
        try:
            results = []
            for query, documents in zip(queries, documents_list):
                reranked = self.rerank_documents(query, documents, top_k, threshold)
                results.append(reranked)
            
            logger.info(f"批量重排序完成，处理了 {len(queries)} 个查询")
            return results
            
        except Exception as e:
            logger.error(f"批量重排序失败: {str(e)}")
            return [docs[:top_k] for docs in documents_list]
    
    def get_model_info(self) -> Dict[str, Any]:
        """获取模型信息"""
        return {
            "model_name": self.model_name,
            "model_loaded": True,
            "model_type": "api-based",
            "description": "使用方舟doubao-seed-1-6-250615模型进行重排序"
        } 